#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
测试图片保留功能的脚本
"""

import os
import sys
from docx import Document
from doc_trans import translate_document, get_document_images_info

def t_image_preservation():
    """测试图片保留功能"""
    
    # 测试文件路径
    input_file = "test.docx"
    output_file = "translated_test_with_images.docx"
    
    # 检查测试文件是否存在
    if not os.path.exists(input_file):
        print(f"测试文件 {input_file} 不存在")
        print("请确保在extract目录中有test.docx文件")
        return False
    
    print("=== 开始测试图片保留功能 ===")
    
    # 1. 获取原始文档的图片信息
    print("1. 获取原始文档图片信息...")
    original_images = get_document_images_info(input_file)
    print(f"原始文档包含 {len(original_images)} 张图片")
    for img_info in original_images:
        print(f"   - 位置: {img_info['position']}, 类型: {img_info['type']}")
    
    # 2. 执行翻译
    print("\n2. 执行文档翻译...")
    try:
        result_file = translate_document(input_file, output_file, 
                                       translate_tables=True,
                                       translation_direction='en_to_zh')
        print(f"翻译完成，输出文件: {result_file}")
    except Exception as e:
        print(f"翻译过程中出错: {e}")
        return False
    
    # 3. 检查翻译后文档的图片信息
    print("\n3. 检查翻译后文档的图片信息...")
    if os.path.exists(output_file):
        translated_images = get_document_images_info(output_file)
        print(f"翻译后文档包含 {len(translated_images)} 张图片")
        for img_info in translated_images:
            print(f"   - 位置: {img_info['position']}, 类型: {img_info['type']}")
        
        # 4. 比较图片数量
        if len(original_images) == len(translated_images):
            print("\n✅ 测试通过：图片数量一致，图片保留成功！")
            return True
        else:
            print(f"\n❌ 测试失败：原始文档有 {len(original_images)} 张图片，翻译后文档有 {len(translated_images)} 张图片")
            return False
    else:
        print("❌ 翻译后文件不存在")
        return False

def extract_and_check_images(docx_path):
    """提取并检查文档中的图片"""
    from doc_trans import extract_images_from_docx
    
    print(f"\n=== 提取 {docx_path} 中的图片 ===")
    
    # 提取图片
    images = extract_images_from_docx(docx_path, "extracted_images")
    print(f"提取到 {len(images)} 张图片:")
    for img in images:
        print(f"  - {img}")
    
    return images

if __name__ == "__main__":
    # 测试图片保留功能
    success = t_image_preservation()
    
    if success:
        # 提取并比较图片
        print("\n=== 详细图片对比 ===")
        
        # 提取原始文档图片
        original_images = extract_and_check_images("test.docx")
        
        # 提取翻译后文档图片
        if os.path.exists("translated_test_with_images.docx"):
            translated_images = extract_and_check_images("translated_test_with_images.docx")
            
            # 比较图片文件
            if len(original_images) == len(translated_images):
                print("\n✅ 图片文件提取成功，数量一致")
            else:
                print(f"\n❌ 图片文件数量不一致: 原始 {len(original_images)} vs 翻译后 {len(translated_images)}")
    
    print("\n=== 测试完成 ===")